/*
 * Will output the RMSE of a prediction for the probe set. Assumes proper format.
 * 
 * JLM: Modified to ignore predictions < 0 (punt). Or, you can pass a fourth
 * parameter that is a list of movie IDs like big.movies.ids (see Makefile), and
 * it will only make predictions for the IDs in that file.
 */
#include <fstream>
#include <set>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

using namespace std;

int main(int argc, char **argv){
  int numValues = 0, lineNumber = 0;
  double sumSquaredValues = 0, prediction, delta, rating;
  double sumPred = 0, sumSqPred = 0, meanPred, variancePred;
  double sumRatings = 0, sumSqRatings = 0, meanRatings, varianceRatings;
  int userId, movieId = -1;
  char probeLine[256], predictionLine[256];
 
  if( argc < 3 ){
    fprintf(stderr, "Need to supply probe ratings file and prediction file for the probe set.\n");
    exit(EXIT_FAILURE);
  }

  set<int> movie_ids;
  if (argc == 4) {
    ifstream is(argv[3]);
    int movie_id;
    while (is >> movie_id) {
      movie_ids.insert(movie_id);
    }
    is.close();
  }

  FILE *probeFile = fopen(argv[1], "r");
  FILE *predictionFile = fopen(argv[2], "r");
  string err_file = string(argv[2]) + ".errors";
  FILE *errorFile = fopen(err_file.c_str(),"w");
  while(fgets(probeLine, 256, probeFile)){//for each probeLine read the corresponding prediction line
    fgets(predictionLine, 256, predictionFile);
    lineNumber++;
    if(probeLine[strlen(probeLine)-2] == ':')//New Movie
      movieId = atoi(probeLine);
    else{ //CustomerID, rating, date // Sum squared difference of rating and prediction
      sscanf(probeLine, "%d, %lf\n", &userId, &rating);
      sscanf(predictionLine, "%lf\n", &prediction);
      // JLM: we allow a prediction of -1 (punt)
      if (prediction > 0
          && (argc != 4 || movie_ids.find(movieId) != movie_ids.end())) {
        if(prediction < 1 || prediction > 5){
          printf("Invalid prediction %lf on line %d\n", prediction, lineNumber);
          exit(EXIT_FAILURE);
        }
        sumPred += prediction;
        sumSqPred += prediction*prediction;
        sumRatings += rating;
        sumSqRatings += rating*rating;
        delta = rating - prediction;
        fprintf(errorFile, "%d %d %f\n", movieId, userId, delta);
        sumSquaredValues += delta * delta;
        numValues++;
      }
    }
  }
  printf("%d pairs RMSE: %.5lf\n", numValues, sqrt(sumSquaredValues/(double)numValues));

  meanPred = sumPred/numValues;
  variancePred = sumSqPred/numValues - meanPred*meanPred;

  meanRatings = sumRatings/numValues;
  varianceRatings = sumSqRatings/numValues - meanRatings*meanRatings;

  printf("prediction mean: %f, variance: %f\n",meanPred, variancePred);
  printf("ratings mean: %f, variance: %f\n",meanRatings, varianceRatings);
  fclose(probeFile);
  fclose(predictionFile);
  fclose(errorFile);
}
